package dataset;

import java.util.*;

public class CustomizedDataset extends Dataset {
    public CustomizedDataset() {

    }

    public CustomizedDataset(List<List<String>> raw, Integer strategy) {
        shape[1] = raw.get(0).size();
        switch (strategy) {
            case 0: // 删除
                List<List<String>> tmp = new ArrayList<>();
                for (int i = 0; i < raw.size(); i++) {
                    boolean skip = false;
                    for (int j = 0; j < shape[1]; j++) {
                        if (raw.get(i).get(j).equals("") || raw.get(i).get(j) == null) {
                            skip = true;
                        }
                    }
                    if (!skip)
                        tmp.add(raw.get(i));
                }
                shape[0] = tmp.size();
                dataset = new double[tmp.size()][shape[1]];
                for (int i = 0; i < tmp.size(); i++) {
                    for (int j = 0; j < shape[1]; j++) {
                        dataset[i][j] = Double.parseDouble(tmp.get(i).get(j));
                    }
                }
                break;
            case 1: // 中位数
                List<List<Double>> mid = new ArrayList<>();
                for (int i = 0; i < shape[1]; i++) {
                    mid.add(new ArrayList<>());
                }
                shape[0] = raw.size();
                dataset = new double[raw.size()][shape[1]];
                List<List<Integer>> recordMid = new ArrayList<>();
                for (int i = 0; i < raw.size(); i++) {
                    for (int j = 0; j < shape[1]; j++) {
                        if (raw.get(i).get(j).equals("") || raw.get(i).get(j) == null) {
                            List<Integer> re = new ArrayList<>();
                            re.add(i);
                            re.add(j);
                            recordMid.add(re);
                        } else {
                            dataset[i][j] = Double.parseDouble(raw.get(i).get(j));
                            mid.get(j).add(dataset[i][j]);
                        }
                    }
                }
                for (int i = 0; i < mid.size(); i++)
                    Collections.sort(mid.get(i));
                for (List<Integer> list : recordMid)
                    dataset[list.get(0)][list.get(1)] = mid.get(list.get(1)).get(mid.get(list.get(1)).size() / 2);
                break;
            case 2: // 平均数
                double[] sum = new double[shape[1]];
                int[] count = new int[shape[1]];
                shape[0] = raw.size();
                dataset = new double[raw.size()][shape[1]];
                List<List<Integer>> recordMean = new ArrayList<>();
                for (int i = 0; i < raw.size(); i++) {
                    for (int j = 0; j < shape[1]; j++) {
                        if (raw.get(i).get(j).equals("") || raw.get(i).get(j) == null) {
                            List<Integer> re = new ArrayList<>();
                            re.add(i);
                            re.add(j);
                            recordMean.add(re);
                        } else {
                            dataset[i][j] = Double.parseDouble(raw.get(i).get(j));
                            sum[j] += dataset[i][j];
                            count[j]++;
                        }
                    }
                }
                for (List<Integer> list : recordMean)
                    dataset[list.get(0)][list.get(1)] = sum[list.get(1)] / count[list.get(1)];
                break;
        }
    }

    public CustomizedDataset(List<List<String>> raw, Integer labelCol, Integer strategy) {
        shape[1] = raw.get(0).size() - 1;
        Map<String, Integer> labelMap = new HashMap<>();
        int cnt = 0;
        List<List<String>> secondRaw = new ArrayList<>();
        List<Integer> labelList = new ArrayList<>();
        for (List<String> line : raw) {
            String col = line.get(labelCol);
            if (col.equals("") || col == null) {
                continue;
            }
            line.remove(col);
            secondRaw.add(line);
            if (!labelMap.containsKey(col))
                labelMap.put(col, cnt++);
            labelList.add(labelMap.get(col));
        }
        switch (strategy) {
            case 0: // 删除
                List<List<String>> tmp = new ArrayList<>();
                List<Integer> labelTmp = new ArrayList<>();
                for (int i = 0; i < secondRaw.size(); i++) {
                    boolean skip = false;
                    for (int j = 0; j < shape[1]; j++) {
                        if (secondRaw.get(i).get(j).equals("") || raw.get(i).get(j) == null) {
                            skip = true;
                        }
                    }
                    if (!skip) {
                        tmp.add(secondRaw.get(i));
                        labelTmp.add(labelList.get(i));
                    }
                }
                shape[0] = tmp.size();
                label = new int[shape[0]];
                dataset = new double[tmp.size()][shape[1]];
                for (int i = 0; i < tmp.size(); i++) {
                    label[i] = labelTmp.get(i);
                    for (int j = 0; j < shape[1]; j++) {
                        dataset[i][j] = Double.parseDouble(tmp.get(i).get(j));
                    }
                }
                break;
            case 1: // 中位数
                List<List<Double>> mid = new ArrayList<>();
                for (int i = 0; i < shape[1]; i++) {
                    mid.add(new ArrayList<>());
                }
                shape[0] = secondRaw.size();
                label = new int[shape[0]];
                dataset = new double[secondRaw.size()][shape[1]];
                List<List<Integer>> recordMid = new ArrayList<>();
                for (int i = 0; i < secondRaw.size(); i++) {
                    label[i] = labelList.get(i);
                    for (int j = 0; j < shape[1]; j++) {
                        if (secondRaw.get(i).get(j).equals("") || secondRaw.get(i).get(j) == null) {
                            List<Integer> re = new ArrayList<>();
                            re.add(i);
                            re.add(j);
                            recordMid.add(re);
                        } else {
                            dataset[i][j] = Double.parseDouble(secondRaw.get(i).get(j));
                            mid.get(j).add(dataset[i][j]);
                        }
                    }
                }
                for (int i = 0; i < mid.size(); i++)
                    Collections.sort(mid.get(i));
                System.out.println(recordMid);
                for (List<Integer> list : recordMid) {
                    dataset[list.get(0)][list.get(1)] = mid.get(list.get(1)).get(mid.get(list.get(1)).size() / 2);
                }
                break;
            case 2: // 平均数
                double[] sum = new double[shape[1]];
                int[] count = new int[shape[1]];
                shape[0] = secondRaw.size();
                label = new int[shape[0]];
                dataset = new double[secondRaw.size()][shape[1]];
                List<List<Integer>> recordMean = new ArrayList<>();
                for (int i = 0; i < secondRaw.size(); i++) {
                    label[i] = labelList.get(i);
                    for (int j = 0; j < shape[1]; j++) {
                        if (secondRaw.get(i).get(j).equals("") || secondRaw.get(i).get(j) == null) {
                            List<Integer> re = new ArrayList<>();
                            re.add(i);
                            re.add(j);
                            recordMean.add(re);
                        } else {
                            dataset[i][j] = Double.parseDouble(secondRaw.get(i).get(j));
                            sum[j] += dataset[i][j];
                            count[j]++;
                        }
                    }
                }
                for (List<Integer> list : recordMean)
                    dataset[list.get(0)][list.get(1)] = sum[list.get(1)] / count[list.get(1)];
                break;
        }
    }

    public Map<Integer, Integer> labelNums() {
        Map<Integer, Integer> map = new HashMap<>();
        for (int i : label)
            map.put(i, map.getOrDefault(i, 0) + 1);
        return map;
    }
}
